﻿using System;

using System.Linq;
using System.Numerics;
using Algoritmia.DataStructures.Lists;
using Algoritmia.Utils;

namespace Algoritmia.Problems.Fourier
{
    public class IterativeFastFourierTransform : IDiscreteFourierTransform
    {
        private readonly int _n;
        private readonly IList<Complex> _w;
        private readonly int[] _bitrev;

        public IterativeFastFourierTransform(int N)
        {
            _n = N;
            _w = Enumerable.Repeat(new Complex(0, 0), N+1).ToAlgoritmiaList();
            int i = 1;
            while (i <= N)
            {
                _w[i] = Complex.Exp(new Complex(0, 2*Math.PI)/i);
                i *= 2;
            }
            int rev = 0;
            _bitrev = new int[N];
            for(i = 0; i < N-1; i++)
            {
                var mask = N / 2;
                _bitrev[i] = rev;
                while ( rev >= mask)
                {
                    rev -= mask;
                    mask >>= 1;
                }
                rev += mask;
            }
            _bitrev[N - 1] = N - 1;
        }

        #region Implementation of IDiscreteFourierTransform

        public IList<Complex> Transform(IList<Complex> complexs)
        {
            var X = (from k in Enumerable.Range(0, _n)
                     select complexs[_bitrev[k]]).ToAlgoritmiaList();
            int m = 2;
            while (m <= _n)
            {
                for(int j = 0; j < _n; j += m)
                {
                    for(int k = 0; k < m/2; k++)
                    {
                        var w = Complex.Pow(_w[m], -k);
                        var Xjk = X[j + k];
                        X[j + k] = Xjk + w*X[j + k + m/2];
                        X[j + k + m/2] = Xjk - w*X[j + k + m/2];
                    }
                }
                m <<= 1;
            }
            return X;
        }

        #endregion
    }
}
